import os
import cv2
import numpy as np
from commonFunctions import *

color_list = {
    "gray": [
        {
            'lower': np.array([0, 0, 60]),
            'upper': np.array([180, 43, 220])
        }
    ],
    "red": [
        {
            'lower': np.array([0, 43, 46]),
            'upper': np.array([10, 255, 255])
        },
        {
            'lower': np.array([156, 43, 46]),
            'upper': np.array([180, 255, 255])
        }
    ],
}


def to_int(*args):
    _ret = []
    for a in args:
        _ret.append(int(a))
    return _ret


def get_color_mask(param_img, _color='gray'):
    # change to hsv model
    _hsv = cv2.cvtColor(param_img, cv2.COLOR_BGR2HSV)
    # get mask
    _mask = np.zeros_like(to_gray(param_img))
    if _color in color_list.keys():
        for _color_info in color_list[_color]:
            _temp = cv2.inRange(_hsv, _color_info['lower'], _color_info['upper'])
            _mask = cv2.bitwise_or(_mask, _temp)
    return _mask


def get_color_contour(_img, _color="gray", _erosion_flag=True):
    _mask = get_color_mask(_img, _color)
    if _erosion_flag:
        kernel = np.ones((5, 5), np.uint8)
        _mask = cv2.erode(_mask, kernel, iterations=1)
    # cv2.imshow('_mask', _mask)
    _contours, hierarchy = cv2.findContours(_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    def cnt_area(cnt):
        area = cv2.contourArea(cnt)
        return area

    _contours.sort(key=cnt_area, reverse=True)
    if len(_contours) > 0:
        return _contours[0], _mask
    else:
        print("None")
        return None, _mask


# 获取白板的中心点
def get_whiteboard_center_point(_img):
    _contour, _mask = get_color_contour(_img)
    cv2.imwrite("whiteboard_mask.jpg", _mask)
    _img_h, _img_w = get_HW(_img)
    if isinstance(_contour, type(None)):
        return _img_w // 2, _img_h // 2
    # (x, y, w, h) = cv2.boundingRect(_contour)
    ((_cx, _cy), (_width, _height), theta) = cv2.minAreaRect(_contour)
    cv2.drawContours(_img, [_contour], -1, (255, 255, 255), 2)
    _cx, _cy, _height, _height = to_int(_cx, _cy, _width, _height)
    return _cx, _cy


if __name__ == '__main__':
    root = r"E:\BaiduDisk\500"
    names = os.listdir(root)
    for name in names:
        img_path = r"E:\BaiduDisk\500\\" + name
        # img_path = r"E:\BaiduDisk\500\IMG_20210918_235051.jpg"
        img = cv2.imread(img_path)
        img = to_resize(img, 4)
        center_x, center_y = get_whiteboard_center_point(img)
        cv2.circle(img, (center_x, center_y), 2, (222, 0, 222), 15)
        cv2.imshow('img', img)
        cv2.waitKey(0)
    cv2.destroyAllWindows()
